import json
import argparse
from trainer import train
from utils.toolkit import NamespaceDict


def merge_config(args, config):
    """merge argparse args to config, replace the value in config if in args"""
    # keys = config.keys()
    for key, value in vars(args).items():
        if value is not None:
            config[key] = value

    return config


def main():
    args = setup_parser().parse_args()
    param = load_json(args.config)
    args = merge_config(args, param)

    # TODO new arg type
    args = NamespaceDict(**args)

    train(args)


def load_json(setting_path):
    with open(setting_path) as data_file:
        param = json.load(data_file)
    return param


def setup_parser():
    parser = argparse.ArgumentParser(
        description="Reproduce of multiple pre-trained incremental learning algorthms."
    )
    parser.add_argument(
        "--config",
        type=str,
        default="./exps/simplecil.json",
        help="Json file of settings.",
    )
    parser.add_argument(
        "--prefix",
        type=str,
        default=None,
        help="Json file of settings.",
    )

    # Incremental learning configs
    parser.add_argument(
        "--init_cls",
        type=int,
        default=None,
        help="init classes",
    )
    parser.add_argument(
        "--inc_cls",
        type=int,
        default=None,
        help="incremental classes",
    )

    # fc config
    parser.add_argument(
        "--fc_temperture",
        action="store_true",
        help="fc_temperture",
    )

    # training config
    parser.add_argument(
        "--interval",
        type=int,
        default=None,
        help="evaluation intervals",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=None,
        help="batch size",
    )
    parser.add_argument(
        "--init_epochs",
        type=int,
        default=None,
        help="first session epoch",
    )
    parser.add_argument(
        "--init_lr",
        type=float,
        default=None,
        help="first session learning rat",
    )
    parser.add_argument(
        "--inc_epochs",
        type=int,
        default=None,
        help="incremental session epoch",
    )
    parser.add_argument(
        "--inc_lr",
        type=float,
        default=None,
        help="incremental session learning rat",
    )
    parser.add_argument(
        "--ca_epochs",
        type=int,
        default=None,
        help="classifier alignment epoch",
    )
    parser.add_argument(
        "--device",
        type=int,
        nargs="+",
        help="devices, accepted in list",
    )
    parser.add_argument(
        "--train_seed",
        type=int,
        default=None,
        help="training seed",
    )
    parser.add_argument(
        "--seed",
        type=int,
        nargs="+",
        default=None,
        help="random seed",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="show GFLOPs",
    )
    parser.add_argument(
        "--early_stop",
        action="store_true",
        help="store the best model",
    )

    # loss config
    parser.add_argument(
        "--distill_alpha",
        type=float,
        default=None,
        help="distill_alpha",
    )
    parser.add_argument(
        "--scale",
        type=float,
        default=None,
        help="distill_alpha",
    )
    parser.add_argument(
        "--margin",
        type=float,
        default=None,
        help="distill_alpha",
    )
    parser.add_argument(
        "--projector",
        type=str,
        default=None,
        help="projector for contrastive learning",
    )

    # model config
    parser.add_argument(
        "--ffn_rank",
        type=int,
        default=None,
        help="Adapter Rank",
    )
    parser.add_argument(
        "--topk",
        type=int,
        default=None,
        help="top k adapter/lora/components",
    )
    parser.add_argument(
        "--align_alpha",
        type=float,
        default=None,
        help="Adapter Rank",
    )

    return parser


if __name__ == "__main__":
    main()
